import numpy as np
import matplotlib.pyplot as plt

from sklearn import datasets
from sklearn.model_selection import train_test_split

from notebook.dnn_utils import initialize_parameters_deep, model_forward, compute_cost, model_backward, \
    update_parameters, predict
from notebook.plot_utils import plot_cost_accuracy, plot_decision_boundary


def build_model(X_train, y_train, X_val, y_val, layers_dims,
                learning_rate=0.001, epochs=100, activation='tanh', print_cost=False):
    # keep track of cost
    costs = []
    train_accs = []
    val_accs = []
    # Parameters initialization.
    parameters = initialize_parameters_deep(layers_dims)
    # Loop (gradient descent)
    for i in range(epochs):
        # Forward propagation: [LINEAR -> RELU]*(L-1) -> LINEAR -> SIGMOID.
        AL, caches = model_forward(X_train, parameters, activation)

        # Compute cost.
        cost = compute_cost(AL, y_train)

        # Backward propagation.
        grads = model_backward(AL, y_train, caches, activation)

        # Update parameters.
        parameters = update_parameters(parameters, grads, learning_rate)

        # Record cost, accuracy
        y_train_pred = predict(X_train, parameters, activation)
        train_acc = np.sum((y_train_pred == y_train) / X_train.shape[1])
        y_val_pred = predict(X_val, parameters, activation)
        val_acc = np.sum((y_val_pred == y_val) / X_val.shape[1])
        costs.append(cost)
        train_accs.append(train_acc)
        val_accs.append(val_acc)
        # Print the cost every 10 training example
        if print_cost and (i + 1) % 10 == 0:
            template = "Cost after iteration {}: {}, Accuracy on train data: {}, Accuracy on validation data: {}"
            print(template.format(i + 1, cost, train_acc, val_acc))

    return parameters, costs, train_accs, val_accs


def main():
    X, y = datasets.make_moons(100000, noise=0.20)
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
    X_val, X_test, y_val, y_test = train_test_split(X_test, y_test, test_size=0.5, random_state=42)

    X_train = X_train.T
    X_val = X_val.T
    X_test = X_test.T
    y_train = y_train.reshape((1, y_train.shape[0]))
    y_val = y_val.reshape((1, y_val.shape[0]))
    y_test = y_test.reshape((1, y_test.shape[0]))
    layers_dims = [2, 3, 1]
    parameters, costs, train_accs, val_accs = build_model(X_train, y_train, X_val, y_val, layers_dims, epochs=100,
                                                          print_cost=True)
    X_plot, y_plot = X_test[:, :1000].copy(), y_test[:, :1000].copy()
    plot_decision_boundary(lambda x: predict(x.T, parameters), X_plot, y_plot)
    plt.show()


if __name__ == '__main__':
    main()
